
// XOR256Block.cpp
#include "bOS.h"
#include "bOSXOR256Block.h"
#include "bOSSHA.h"
#include "bOSDoubleBuffering.h"

#include <exception>
#include <strstream>

using namespace std;

char const* CXOR256Block::sm_szErrorMsgXB1 = "Illegal External Rounds!";
char const* CXOR256Block::sm_szErrorMsgXB2 = "Illegal Internal Rounds!";

//Null chain
char const CXOR256Block::sm_chain0[BLOCK_MAX] = {0};

//CONSTRUCTOR
CXOR256Block::CXOR256Block()
{
	m_blockSize = -1;
	m_keylength = KEY_MAX;
}

//DESTRUCTOR
CXOR256Block::~CXOR256Block()
{
	if(m_pucXOR != NULL)
		delete [] m_pucXOR;
	if(m_puc256 != NULL)
		delete [] m_puc256;
};

void CXOR256Block::Initialize(char const* keydata, int keydatalength, char const* chain,
	int blockSize, int iORounds, int iIRounds, int iMode, int iPadding)
{
	//Check Initialization Data
	if(NULL == keydata)
		throw runtime_error(string(sm_szErrorMsg4));
	if(keydatalength<1)
		throw runtime_error(string(sm_szErrorMsg5));
	if(blockSize<1 || blockSize>BLOCK_MAX)
		throw runtime_error(string(sm_szErrorMsg6));
	if(iORounds<1)
		throw runtime_error(string(sm_szErrorMsgXB1));
	if(iIRounds<1)
		throw runtime_error(string(sm_szErrorMsgXB2));
	if(iMode<ECB || iMode>CFB)
		throw runtime_error(string(sm_szErrorMsg2));
	if(iPadding<ZEROES || iPadding>PKCS7)
		throw runtime_error(string(sm_szErrorMsg3));
	m_iORounds = iORounds;
	m_iORounds1 = m_iORounds-1;
	m_iIRounds = iIRounds;
	m_iIRounds1 = m_iIRounds-1;
	m_iMode = iMode;
	m_iPadding = iPadding;
	//Create the Key from Key Data
	int i, j;
	char key[KEY_MAX];
	for(i=0,j=0; i<m_keylength; i++,j=(j+1)%keydatalength)
		key[i] = keydata[j];
	bool bSameKey = false;
	bool bSameChain = false;
	if(true == m_bInit)
	{
		//Only if already initialized
		if(m_blockSize == blockSize)
		{
			//Check the Chain if is the same
			if(0 == memcmp(m_apchain0.get(), chain, m_blockSize))
				bSameChain = true;
		}
		//Check the Key if is the same
		if(0 == memcmp(m_apKey.get(), key, m_keylength))
			bSameKey = true;
	}
	if(true == bSameChain)
		//Just Reset
		memcpy(m_apchain.get(), m_apchain0.get(), m_blockSize);
	else
	{
		if(m_blockSize != blockSize)
		{
			m_blockSize = blockSize;
			m_iBlockSize1 = m_blockSize-1;
			m_iProd = m_blockSize * m_iIRounds;
			m_iDelta = m_iProd*m_iORounds1;
			//Initialize the chain
			m_apchain0 = auto_ptr<char>(new char[m_blockSize]);
			m_apchain = auto_ptr<char>(new char[m_blockSize]);
			//Initializing temporary buffer
			m_apTemp = auto_ptr<unsigned char>(new unsigned char[m_blockSize]);
		}
		memcpy(m_apchain0.get(), chain, m_blockSize);
		memcpy(m_apchain.get(), chain, m_blockSize);
	}
	if(true == bSameKey)
	{
		//Fast Initialization
		m_oArcfourPRNG.Reset();
		return;
	}
	m_apKey = auto_ptr<char>(new char[m_keylength]);
	memcpy(m_apKey.get(), key, m_keylength);
	//Calculating the constants
	m_oArcfourPRNG.SetKey((unsigned char*)key, m_keylength);
	//The number of constants is given by ORounds * IRounds * BlockSize
	int iSize = m_iORounds*m_iIRounds*m_blockSize;
	m_pucXOR = new unsigned char[iSize];
	m_puc256 = new unsigned char[iSize];
	unsigned char* pucXOR = &m_pucXOR[0];
	unsigned char* puc256 = &m_puc256[0];
	for(i=0; i<iSize; i++)
	{
		*(pucXOR++) = m_oArcfourPRNG.Rand();
		*(puc256++) = m_oArcfourPRNG.Rand();
	}
	//Initialization Flag
	m_bInit = true;
}

void CXOR256Block::ResetChain()
{
	if(false==m_bInit)
		throw runtime_error(string(sm_szErrorMsg1));
	memcpy(m_apchain.get(), m_apchain0.get(), m_blockSize);
	//Is including ArcFourPRNG resetting
	m_oArcfourPRNG.Reset();
}

//Compute Signature
void CXOR256Block::Signature(char* pcSig)
{
	//11+256+3+3+3+1+1+1
	char acSigData[280] = {0};
	strcat(acSigData, "XOR256BLOCK");
	int iLen = strlen(acSigData);
	memcpy(acSigData+iLen, m_apKey.get(), m_keylength);
	sprintf(acSigData+iLen+m_keylength, "%d%d%d%d%d", m_blockSize, m_iIRounds, m_iORounds, m_iMode, m_iPadding);
	CSHA oSHA;
	oSHA.AddData(acSigData, strlen(acSigData));
	oSHA.FinalDigest(pcSig);
}

void CXOR256Block::EncryptDirect(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	for(int i=0; i<m_blockSize; i++,pucBlock++)
	{
		//The First Internal Round
		(*pucBlock ^= ucPrev^*(pucXOR++)) += *(puc256++);
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=1; j<m_iIRounds; j++)
			(*pucBlock ^= *(pucXOR++)) += *(puc256++);
		ucPrev = *pucBlock;
	}
}

void CXOR256Block::EncryptReverse(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	pucBlock += m_iBlockSize1;
	for(int i=0; i<m_blockSize; i++,pucBlock--)
	{
		//The First Internal Round
		(*pucBlock ^= ucPrev^*(pucXOR++)) += *(puc256++);
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=1; j<m_iIRounds; j++)
			(*pucBlock ^= *(pucXOR++)) += *(puc256++);
		ucPrev = *pucBlock;
	}
}

void CXOR256Block::EncryptDirect1(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	for(int i=0; i<m_blockSize; i++,pucBlock++)
	{
		ucPrev ^= *pucBlock;
		//The First Internal Round
		(*pucBlock = ucPrev^*(pucXOR++)) += *(puc256++);
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=1; j<m_iIRounds; j++)
			(*pucBlock ^= *(pucXOR++)) += *(puc256++);
	}
}

void CXOR256Block::EncryptReverse1(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	pucBlock += m_iBlockSize1;
	for(int i=0; i<m_blockSize; i++,pucBlock--)
	{
		ucPrev ^= *pucBlock;
		//The First Internal Round
		(*pucBlock = ucPrev^*(pucXOR++)) += *(puc256++);
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=1; j<m_iIRounds; j++)
			(*pucBlock ^= *(pucXOR++)) += *(puc256++);
	}
}

void CXOR256Block::DecryptDirect(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	unsigned char ucTemp;
	unsigned char const* pucXOR_1;
	unsigned char const* puc256_1;
	for(int i=0; i<m_blockSize; i++,pucBlock++)
	{
		pucXOR += m_iIRounds;
		puc256 += m_iIRounds;
		pucXOR_1 = pucXOR - 1;
		puc256_1 = puc256 - 1;
		ucTemp = *pucBlock;
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=m_iIRounds1; j>0; j--,pucXOR_1--,puc256_1--)
		{
			if(*puc256_1 <= *pucBlock)
				*pucBlock -= *puc256_1;
			else
				(*pucBlock += ~(*puc256_1))++;
			*pucBlock ^= *pucXOR_1;
		}
		//The First Internal Round
		if(*puc256_1 <= *pucBlock)
			*pucBlock -= *puc256_1;
		else
			(*pucBlock += ~(*puc256_1))++;
		*pucBlock ^= ucPrev ^ *pucXOR_1;
		ucPrev = ucTemp;
	}
}

void CXOR256Block::DecryptReverse(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	unsigned char ucTemp;
	unsigned char const* pucXOR_1;
	unsigned char const* puc256_1;
	pucBlock += m_iBlockSize1;
	for(int i=0; i<m_blockSize; i++,pucBlock--)
	{
		pucXOR += m_iIRounds;
		puc256 += m_iIRounds;
		pucXOR_1 = pucXOR - 1;
		puc256_1 = puc256 - 1;
		ucTemp = *pucBlock;
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=m_iIRounds1; j>0; j--,pucXOR_1--,puc256_1--)
		{
			if(*puc256_1 <= *pucBlock)
				*pucBlock -= *puc256_1;
			else
				(*pucBlock += ~(*puc256_1))++;
			*pucBlock ^= *pucXOR_1;
		}
		//The First Internal Round
		if(*puc256_1 <= *pucBlock)
			*pucBlock -= *puc256_1;
		else
			(*pucBlock += ~(*puc256_1))++;
		*pucBlock ^= ucPrev ^ *pucXOR_1;
		ucPrev = ucTemp;
	}
}

void CXOR256Block::DecryptDirect1(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	unsigned char const* pucXOR_1;
	unsigned char const* puc256_1;
	for(int i=0; i<m_blockSize; i++,pucBlock++)
	{
		pucXOR += m_iIRounds;
		puc256 += m_iIRounds;
		pucXOR_1 = pucXOR - 1;
		puc256_1 = puc256 - 1;
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=m_iIRounds1; j>0; j--,pucXOR_1--,puc256_1--)
		{
			if(*puc256_1 <= *pucBlock)
				*pucBlock -= *puc256_1;
			else
				(*pucBlock += ~(*puc256_1))++;
			*pucBlock ^= *pucXOR_1;
		}
		//The First Internal Round
		if(*puc256_1 <= *pucBlock)
			*pucBlock -= *puc256_1;
		else
			(*pucBlock += ~(*puc256_1))++;
		*pucBlock ^= ucPrev ^ *pucXOR_1;
		ucPrev ^= *pucBlock;
	}
}

void CXOR256Block::DecryptReverse1(unsigned char* pucBlock, unsigned char const* pucXOR, unsigned char const* puc256)
{
	unsigned char ucPrev = 0;
	unsigned char const* pucXOR_1;
	unsigned char const* puc256_1;
	pucBlock += m_iBlockSize1;
	for(int i=0; i<m_blockSize; i++,pucBlock--)
	{
		pucXOR += m_iIRounds;
		puc256 += m_iIRounds;
		pucXOR_1 = pucXOR - 1;
		puc256_1 = puc256 - 1;
		//The Last m_iIRounds-1 Internal Rounds
		for(int j=m_iIRounds1; j>0; j--,pucXOR_1--,puc256_1--)
		{
			if(*puc256_1 <= *pucBlock)
				*pucBlock -= *puc256_1;
			else
				(*pucBlock += ~(*puc256_1))++;
			*pucBlock ^= *pucXOR_1;
		}
		//The First Internal Round
		if(*puc256_1 <= *pucBlock)
			*pucBlock -= *puc256_1;
		else
			(*pucBlock += ~(*puc256_1))++;
		*pucBlock ^= ucPrev ^ *pucXOR_1;
		ucPrev ^= *pucBlock;
	}
}

void CXOR256Block::EncryptBlock(unsigned char* pucBlock)
{
	unsigned char const* pucXOR = m_pucXOR;
	unsigned char const* puc256 = m_puc256;
	bool bDir = true;
	for(int i=0; i<m_iORounds1; i++,pucXOR+=m_iProd,puc256+=m_iProd)
	{
		if(true == bDir)
			EncryptDirect(pucBlock, pucXOR, puc256);
		else
			EncryptReverse(pucBlock, pucXOR, puc256);
		bDir = !bDir;
	}
	if(true == bDir)
		EncryptDirect1(pucBlock, pucXOR, puc256);
	else
		EncryptReverse1(pucBlock, pucXOR, puc256);
}

void CXOR256Block::DecryptBlock(unsigned char* pucBlock)
{
	unsigned char const* pucXOR = m_pucXOR + m_iDelta;
	unsigned char const* puc256 = m_puc256 + m_iDelta;
	bool bDir;	
	if(1 == m_iORounds%2)
	{
		DecryptDirect1(pucBlock, pucXOR, puc256);
		bDir = false;
	}
	else
	{
		DecryptReverse1(pucBlock, pucXOR, puc256);
		bDir = true;
	}
	pucXOR-=m_iProd;
	puc256-=m_iProd;
	for(int i=0; i<m_iORounds1; i++,pucXOR-=m_iProd,puc256-=m_iProd)
	{
		if(true == bDir)
			DecryptDirect(pucBlock, pucXOR, puc256);
		else
			DecryptReverse(pucBlock, pucXOR, puc256);
		bDir = !bDir;
	}
}

void CXOR256Block::Encrypt(char const* in, char* result, size_t n)
{
	if(false==m_bInit)
		throw runtime_error(string(sm_szErrorMsg1));
	//n should be > 0 and multiple of m_blockSize
	if(n<1 || n%m_blockSize!=0)
		throw runtime_error(string(sm_szErrorMsg6));
	unsigned int ui;
	unsigned char const* pucCurrentIn;
	unsigned char* pucCurrentOut;
	if(m_iMode == CBC) //CBC mode, using the Chain
	{
		for(ui=0,pucCurrentIn=reinterpret_cast<unsigned char const*>(in),
			pucCurrentOut=reinterpret_cast<unsigned char*>(result); ui<n/m_blockSize; ui++)
		{
			memcpy(pucCurrentOut, pucCurrentIn, m_blockSize);
			Xor(reinterpret_cast<char*>(pucCurrentOut), m_apchain.get());
			EncryptBlock(pucCurrentOut);
			memcpy(m_apchain.get(), pucCurrentOut, m_blockSize);
			pucCurrentIn += m_blockSize;
			pucCurrentOut += m_blockSize;
		}
	}
	else if(m_iMode == CFB) //CFB mode, using the Chain
	{
		for(ui=0,pucCurrentIn=reinterpret_cast<unsigned char const*>(in),
			pucCurrentOut=reinterpret_cast<unsigned char*>(result); ui<n/m_blockSize; ui++)
		{
			EncryptBlock(reinterpret_cast<unsigned char*>(m_apchain.get()));
			memcpy(pucCurrentOut, pucCurrentIn, m_blockSize);
			Xor(reinterpret_cast<char*>(pucCurrentOut), m_apchain.get());
			memcpy(m_apchain.get(), pucCurrentOut, m_blockSize);
			pucCurrentIn += m_blockSize;
			pucCurrentOut += m_blockSize;
		}
	}
	else //ECB mode, not using the Chain
	{
		for(ui=0,pucCurrentIn=reinterpret_cast<unsigned char const*>(in),
			pucCurrentOut=reinterpret_cast<unsigned char*>(result); ui<n/m_blockSize; ui++)
		{
			memcpy(pucCurrentOut, pucCurrentIn, m_blockSize);
			EncryptBlock(pucCurrentOut);
			pucCurrentIn += m_blockSize;
			pucCurrentOut += m_blockSize;
		}
	}
}

void CXOR256Block::Decrypt(char const* in, char* result, size_t n)
{
	if(false==m_bInit)
		throw runtime_error(string(sm_szErrorMsg1));
	//n should be > 0 and multiple of m_blockSize
	if(n<1 || n%m_blockSize!=0)
		throw runtime_error(string(sm_szErrorMsg6));
	unsigned int ui;
	unsigned char const* pucCurrentIn;
	unsigned char* pucCurrentOut;
	if(m_iMode == CBC) //CBC mode, using the Chain
	{
		for(ui=0,pucCurrentIn=reinterpret_cast<unsigned char const*>(in),
			pucCurrentOut=reinterpret_cast<unsigned char*>(result); ui<n/m_blockSize; ui++)
		{
			memcpy(pucCurrentOut, pucCurrentIn, m_blockSize);
			memcpy(m_apTemp.get(), pucCurrentOut, m_blockSize);
			DecryptBlock(pucCurrentOut);
			Xor(reinterpret_cast<char*>(pucCurrentOut), m_apchain.get());
			memcpy(m_apchain.get(), m_apTemp.get(), m_blockSize);
			pucCurrentIn += m_blockSize;
			pucCurrentOut += m_blockSize;
		}
	}
	else if(m_iMode == CFB) //CFB mode, using the Chain, not using Decrypt()
	{
		for(ui=0,pucCurrentIn=reinterpret_cast<unsigned char const*>(in),
			pucCurrentOut=reinterpret_cast<unsigned char*>(result); ui<n/m_blockSize; ui++)
		{
			EncryptBlock(reinterpret_cast<unsigned char*>(m_apchain.get()));
			memcpy(pucCurrentOut, pucCurrentIn, m_blockSize);
			memcpy(m_apTemp.get(), pucCurrentOut, m_blockSize);
			Xor(reinterpret_cast<char*>(pucCurrentOut), m_apchain.get());
			memcpy(m_apchain.get(), m_apTemp.get(), m_blockSize);
			pucCurrentIn += m_blockSize;
			pucCurrentOut += m_blockSize;
		}
	}
	else //ECB mode, not using the Chain
	{
		for(ui=0,pucCurrentIn=reinterpret_cast<unsigned char const*>(in),
			pucCurrentOut=reinterpret_cast<unsigned char*>(result); ui<n/m_blockSize; ui++)
		{
			memcpy(pucCurrentOut, pucCurrentIn, m_blockSize);
			DecryptBlock(pucCurrentOut);
			pucCurrentIn += m_blockSize;
			pucCurrentOut += m_blockSize;
		}
	}
}

void CXOR256Block::EncryptFile(string const& rostrFileIn, string const& rostrFileOut)
{
	if(false==m_bInit)
		throw runtime_error(string(sm_szErrorMsg1));
	//Check if the same file for input and output
	if(rostrFileIn == rostrFileOut)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg8 << rostrFileIn << "!" << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Open Input File
	ifstream in(rostrFileIn.c_str(), ios::binary);
	if(!in)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg7 << rostrFileIn << "!" << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Open Output File
	ofstream out(rostrFileOut.c_str(), ios::binary);
	if(!out)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg7 << rostrFileOut << "!" << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Computing the signature
	char acSig[33] = {0};
	Signature(acSig);
	//Writing the Signature
	out.write(acSig, 32);
	//Resetting the chain
	ResetChain();
	//Reading from file
	char szLargeBuff[BUFF_LEN+1] = {0};
	int DATA_LEN = (BUFF_LEN/m_blockSize/2)*m_blockSize;
	char szBuffIn[BUFF_LEN/2+1] = {0};
	char szBuffOut[BUFF_LEN/2+1] = {0};
	CDoubleBuffering oDoubleBuffering(in, szLargeBuff, BUFF_LEN, DATA_LEN);
	int iRead;
	while((iRead=oDoubleBuffering.GetData(szBuffIn)) > 0)
	{
		if(iRead < DATA_LEN)
			iRead = Pad(szBuffIn, iRead);
		//Encrypting
		Encrypt(szBuffIn, szBuffOut, iRead);
		out.write(szBuffOut, iRead);
	}
	in.close();
	out.close();
}

void CXOR256Block::DecryptFile(string const& rostrFileIn, string const& rostrFileOut)
{
	if(false==m_bInit)
		throw runtime_error(string(sm_szErrorMsg1));
	//Check if the same file for input and output
	if(rostrFileIn == rostrFileOut)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg8 << rostrFileIn << "!" << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Open Input File
	ifstream in(rostrFileIn.c_str(), ios::binary);
	if(!in)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg7 << rostrFileIn << "!" << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Open Output File
	ofstream out(rostrFileOut.c_str(), ios::binary);
	if(!out)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg7 << rostrFileOut << "!" << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Computing the signature
	char acSig[33] = {0};
	Signature(acSig);
	char acSig1[33] = {0};
	//Reading the Signature
	in.read(acSig1, 32);
	//Compare the signatures
	if(memcmp(acSig1, acSig, 32) != 0)
	{
		ostrstream ostr;
		ostr << sm_szErrorMsg9 << rostrFileIn << sm_szErrorMsg10 << ends;
		string ostrMsg = ostr.str();
		ostr.freeze(false);
		throw runtime_error(ostrMsg);
	}
	//Resetting the chain
	ResetChain();
	//Reading from file
	char szLargeBuff[BUFF_LEN+1] = {0};
	int DATA_LEN = (BUFF_LEN/m_blockSize/2)*m_blockSize;
	char szBuffIn[BUFF_LEN/2+1] = {0};
	char szBuffOut[BUFF_LEN/2+1] = {0};
	CDoubleBuffering oDoubleBuffering(in, szLargeBuff, BUFF_LEN, DATA_LEN);
	int iRead;
	while((iRead=oDoubleBuffering.GetData(szBuffIn)) > 0)
	{
		//Encrypting
		Decrypt(szBuffIn, szBuffOut, iRead);
		out.write(szBuffOut, iRead);
	}
	in.close();
	out.close();
}

